package org.summer.gateway.security;

import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextImpl;
import org.springframework.security.web.server.context.ServerSecurityContextRepository;
import org.springframework.stereotype.Repository;
import org.springframework.web.server.ServerWebExchange;
import org.summer.gateway.cache.LoginUserCacheHandler;
import org.summer.gateway.jwt.JsonWebToken;
import org.summer.gateway.jwt.JsonWebTokenResolver;
import reactor.core.publisher.Mono;

/**
 * 网关安全上下文存储器
 */
@Log4j2
@Repository
@RequiredArgsConstructor
public class GatewaySecurityContextRepository implements ServerSecurityContextRepository {
    private final JsonWebTokenResolver<JsonWebToken> jwtResolver;
    private final LoginUserCacheHandler cacheHandler;

    @Override
    public Mono<Void> save(ServerWebExchange exchange, SecurityContext context) {
        log.debug("用户登录请求, 保存用户登录上下文");
        return Mono.empty();
    }

    /**
     * 加载用户登录上下文, 如果反回了SecurityContext,权限认证管理器拒绝访问,则调用ServerAccessDeniedHandler.handle
     * 如果范围empty,则调用ServerAuthenticationEntryPoint.commence
     *
     * @param exchange the exchange to look up the {@link SecurityContext}
     * @return 安全上下文
     */
    @Override
    public Mono<SecurityContext> load(ServerWebExchange exchange) {
        ServerHttpRequest request = exchange.getRequest();
        String token = request.getHeaders().getFirst(HttpHeaders.AUTHORIZATION);
        log.debug("用户请求: {} - {} - {}", request.getMethod(), request.getPath(), request.getRemoteAddress());
        JsonWebToken jsonWebToken = jwtResolver.parseToken(token);
        if (jsonWebToken.isError()) {
            log.error("用户请求头中令牌解析失败: {}", jsonWebToken.getErrorMessage());
            exchange.getResponse().setStatusCode(HttpStatus.UNAUTHORIZED);
            return Mono.empty();
        }
        if (jsonWebToken.isExpired()) {
            log.warn("用户请求头中令牌已过期: {}", jsonWebToken.getExpireTime());
            exchange.getResponse().setStatusCode(HttpStatus.FAILED_DEPENDENCY);
            return Mono.empty();
        }
        removeToken(exchange);
        return cacheHandler.get(jsonWebToken.getUsername())
                .switchIfEmpty(Mono.empty())
                .flatMap(user -> {
                    UsernamePasswordAuthenticationToken authenticationToken = new UsernamePasswordAuthenticationToken(user.getUsername(), user.getTokenId(), user.getAuthorities());
                    return Mono.just(new SecurityContextImpl(authenticationToken));
                });
    }

    /**
     * 移除用户请求中的token信息
     */
    private void removeToken(ServerWebExchange exchange){
        exchange.mutate().request(exchange
                .getRequest()
                .mutate()
                .headers(httpHeaders -> httpHeaders.remove(HttpHeaders.AUTHORIZATION))
                .build()).build();
    }
}
